package com.ng8866.gateway.balancer;

import cn.hutool.core.util.RandomUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.DefaultResponse;
import org.springframework.cloud.client.loadbalancer.EmptyResponse;
import org.springframework.cloud.client.loadbalancer.Request;
import org.springframework.cloud.client.loadbalancer.Response;
import org.springframework.cloud.loadbalancer.core.NoopServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
import org.springframework.cloud.loadbalancer.core.SelectedInstanceCallback;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 自定义负载均衡
 */
@Slf4j
public class LeastConnectionLoadBalancer implements ReactorServiceInstanceLoadBalancer {
    private ObjectProvider<ServiceInstanceListSupplier> provider;
    // 存储服务、路由及路由当前连接数的 map
    public static ConcurrentHashMap<String, ConcurrentHashMap<String, AtomicInteger>> loadMap = new ConcurrentHashMap<>();
    public static ConcurrentHashMap<String, AtomicInteger> countMap = new ConcurrentHashMap<>();
    public static Boolean isCount = false;

    /**
     * 设置是否开启计数功能
     *
     * @param count 是否开启计数功能的布尔值，true表示开启，false表示关闭
     */
    public static void setCount(Boolean count) {
        isCount = count;
        if (!isCount) {
            countMap.clear();
        }
    }

    /**
     * 设置服务ID对应的计数器值
     *
     * @param serviceId 服务ID
     *                  如果isCount为false，则直接返回，不进行任何操作
     *                  如果服务ID对应的计数器不存在，则在countMap中添加一个新的计数器，并初始化为0
     *                  如果服务ID对应的计数器已经存在，则将其值自增1
     */
    public static void setCountMap(String serviceId) {
        if (!isCount) {
            return;
        }
        if (countMap.get(serviceId) == null) {
            AtomicInteger count = new AtomicInteger();
            countMap.putIfAbsent(serviceId, count);
        } else {
            countMap.get(serviceId).incrementAndGet();
        }
        // 计数器值达到最大值，则将关闭计数
        if (countMap.get(serviceId).get() == Integer.MAX_VALUE - 1000) {
            setCount(false);
        }
    }

    // 构造方法
    public LeastConnectionLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> provider) {
        log.debug("==>进入自定义负载均衡 构造方法");
        this.provider = provider;
    }

    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {
        ServiceInstanceListSupplier supplier = provider.getIfAvailable(NoopServiceInstanceListSupplier::new);
        return supplier.get(request).next().map(serviceInstances -> processInstanceResponse(supplier, serviceInstances));
    }

    /**
     * 处理服务实例响应
     *
     * @param supplier         服务实例列表提供者
     * @param serviceInstances 服务实例列表
     * @return 处理后的服务实例响应
     */
    private Response<ServiceInstance> processInstanceResponse(ServiceInstanceListSupplier supplier, List<ServiceInstance> serviceInstances) {
        Response<ServiceInstance> serviceInstanceResponse = getInstanceResponse(serviceInstances);
        if (serviceInstanceResponse.hasServer() && supplier != null) {
            if (supplier instanceof SelectedInstanceCallback) {
                ((SelectedInstanceCallback) supplier).selectedServiceInstance(serviceInstanceResponse.getServer());
            } else {
                log.debug("==>进入自定义负载均衡, supplier不是预期的SelectedInstanceCallback类型{}", supplier);
            }
        }
        return serviceInstanceResponse;
    }


    private Response<ServiceInstance> getInstanceResponse(List<ServiceInstance> instances) {
        if (instances == null || instances.isEmpty()) {
            log.info("==>进入自定义负载均衡, 没有可用服务");
            return new EmptyResponse();
        }
        ServiceInstance firstInstance = instances.get(0);
        if (firstInstance == null) {
            log.info("==>进入自定义负载均衡, 首个实例为空");
            return new EmptyResponse();
        }
        String serviceId = firstInstance.getServiceId();
        log.info("==>进入自定义负载均衡, 使用服务：{}", serviceId);
        Map<String, AtomicInteger> serviceMap = loadMap.get(serviceId);
        if (serviceMap == null || serviceMap.isEmpty()) {
            ServiceInstance randomInstance = RandomUtil.randomEle(instances);
            log.info("==>进入自定义负载均衡, 服务映射为空, 随机负载：{}", randomInstance);
            return new DefaultResponse(randomInstance);
        }
        if (instances.size() == 1) {
            log.info("==>进入自定义负载均衡, 唯一负载：{}", getInstanceKey(firstInstance.getHost(), firstInstance.getPort()));
            serviceMap.getOrDefault(getInstanceKey(firstInstance.getHost(), firstInstance.getPort()), new AtomicInteger())
                    .incrementAndGet();
            return new DefaultResponse(firstInstance);
        }
        String instanceId = getMinAndIncrementKey(serviceMap);
        ServiceInstance chosenInstance = getServiceInstance(instanceId, instances);
        if (chosenInstance == null) {
            ServiceInstance fallbackInstance = RandomUtil.randomEle(instances);
            log.info("==>进入自定义负载均衡, 获取不到服务实例 {}, 随机负载：{}", instanceId, fallbackInstance);
            return new DefaultResponse(fallbackInstance);
        }
        log.info("==>进入自定义负载均衡, 选择负载：{}", instanceId);
        return new DefaultResponse(chosenInstance);
    }


    /**
     * 获取服务
     */
    private ServiceInstance getServiceInstance(String instanceId, List<ServiceInstance> instances) {
        for (ServiceInstance instance : instances) {
            if (Objects.equals(instanceId, getInstanceKey(instance.getHost(), instance.getPort()))) {
                return instance;
            }
        }
        return null;
    }

    /**
     * 获取具有最小值的键，并增加该键对应的计数
     */
    private synchronized String getMinAndIncrementKey(Map<String, AtomicInteger> map) {
        if (map == null || map.isEmpty()) {
            return null; // 防止空指针异常，当map为空或没有元素时直接返回null
        }
        log.debug("前==>{}", map);
        int minValue = Integer.MAX_VALUE;
        String minKey = null;
        // 获取值最小的key
        for (Map.Entry<String, AtomicInteger> entry : map.entrySet()) {
            int currentValue = entry.getValue().get();
            if (currentValue < minValue) {
                minValue = currentValue;
                minKey = entry.getKey();
            }
        }
        // 值自增
        if (minKey != null) {
            AtomicInteger count = map.get(minKey);
            if (count != null) {
                count.incrementAndGet();
            }
        }
        log.debug("后==>{}", map);
        return minKey;
    }


    public static String getInstanceKey(String ip, int port) {
        return ip + ":" + port;
    }
}
